This PR contains three independent, bit-exact performance fixes to the DE /
anndata metric pipeline. All leave outputs numerically identical; none requires
re-baselining metrics. They were found in sequence: each fix cleared the
reigning bottleneck on a ~18k-perturbation screen and exposed the next.

1. Vectorize discrimination_score (anndata-pair metric) — closes discrimination_score recomputes an n_pert x n_pert distance matrix one row at a time #237.
2. Memoize the DE overlap rank-matrix pivot and set-ify column membership in
   compute_overlap / get_top_genes (DE metric).
3. Replace the per-perturbation full-table scans in DENsigCounts and
   compute_generic_auc (pr/roc) with a single grouped/partitioned slice
   (DE metrics).

Fix 1 — Vectorize discrimination_score (closes #237)

What

discrimination_score looped over n_pert perturbations, calling
pairwise_distances once per perturbation to compute a single row of an
n_pert x n_pert distance matrix. This computes the full matrix in one call
and ranks each perturbation by locating its column's position in the per-row
sorted order.

The target-gene-exclusion path (default for expression data) drops a different
feature column per perturbation, so a single unmasked call can't reproduce it.
The full matrix is computed once and corrected per row with an exact,
vectorized rank-1 update that removes the target gene's contribution:

• l1: subtract |pred_g - real_g|
• l2: sqrt(d^2 - (pred_g - real_g)^2)
• cosine: drop the column from the dot product and both norms (masked
  squared norms are clipped at 0 so a target-gene-dominated effect can't round
  negative into a NaN)

Metrics without a closed-form correction fall back to exact per-row masked
distances, and duplicate gene names matching a perturbation are handled by an
exact per-row safety net.

Ranking uses np.where(order == arange[:, None]) to find each row's matching
column rather than a second argsort over the full matrix — identical result,
but it avoids an extra O(n_pert^2) int64 array and a second sort (thanks to
the review suggestion). At n_pert=10000 the ranking step alone drops from
~9.8s / 2.4 GB to ~5.7s / 1.7 GB, a large share of the l2/cosine runtime.

Parity (no behavior change)

Output is numerically identical to the original loop. Across 216 synthetic
configurations (metric x exclude_target_gene x embed_key x seed x
targeting-fraction) the normalized ranks match bit-for-bit (worst
|delta-rank| = 0), and the equivalence is covered by the unit tests.

Benchmark

Apple M2 Pro, Python 3.12.12, numpy 2.4.6 / scipy 1.17.1 / scikit-learn 1.8.0,
n_genes=2000 (old = original per-perturbation loop, new = vectorized):

(n_pert=10000 old times are from the same reference loop on the same machine;
the loop is unchanged by this PR.)

l2/cosine use the BLAS dot-product trick, so one matrix multiply replaces
n_pert dispatched calls and the speedup grows with n_pert; l1 (manhattan)
is a non-BLAS kernel, so its gain is the removed per-call dispatch overhead.
Trade-off: the full matrix is O(n_pert^2) memory vs O(n_pert) per iteration
before — comfortable for typical screen sizes.

Reproduce: python benchmarks/bench_discrimination_score.py

Fix 2 — DE overlap metric: memoize the rank-matrix pivot + O(1) column membership

cell_eval/_types/_de.py — DEComparison.compute_overlap / DEResults.get_top_genes.

Motivation

An OmniPert report_only validate-only run over a HuangChu context with
~18,000 perturbations (16 CPU, 24h wall) ran for 8.8 hours and never
reached results.csv. Faulthandler sampling of the live process (524 periodic
stack dumps) put 513/524 ≈ 98% of all samples inside a single metric:

de_overlap_metric -> DEComparison.compute_overlap -> DEResults.get_top_genes
  hottest leaf: polars frame.py ... in `columns`

The DE compute itself (pdex) was already fast (~135s); this is purely the
metric pass, and it ate the entire budget before the rest of the DE metrics
and the anndata-pair metrics even started.

Two compounding inefficiencies

1. Redundant rank-matrix pivots (the big one). compute_overlap calls
self.real.get_top_genes(sort_by, fdr_threshold) and self.pred.get_top_genes(...)
at the top of every invocation. get_top_genes builds a polars .pivot() with
one column per perturbation (~18k-wide). The de/full profile registers 10
variants of this metric — metrics/_impl.py does
for metric in ["overlap","precision"]: for n in [None,50,100,200,500] with
kwargs={"k": n, "metric": metric} and the default sort_by / fdr_threshold.
So the identical pair of 18k-column pivots is rebuilt 10 times; k only
affects the per-pert truncation genes[:k_eff] inside the loop, never the
matrix.

Fix: memoize get_top_genes keyed by (sort_by, fdr_threshold) on the
per-side DEResults (a dataclass field excluded from init/repr/eq).
DESortBy is an enum (hashable), so the key is safe.

2. O(perts²) column membership. The per-pert loop did
if pert not in real_sig_rank_matrix.columns or pert not in pred_sig_rank_matrix.columns:.
polars rebuilds a fresh ~18k-element list on every .columns access, twice
per pert, over ~18k perts (~3×10⁸ list rebuilds — this is the columns leaf the
sampler kept hitting).

Fix: precompute real_cols = set(real_sig_rank_matrix.columns) and
pred_cols = set(...) once before the loop and test against the sets. (Column
selection matrix[pert] is unchanged.)

Impact

The overlap pass drops from 10 redundant wide pivots + an O(perts²) membership
loop to 2 pivots + an O(perts) loop. What remains irreducible is one pivot
per side plus the per-pert intersect; everything the fix removes scales with
perturbation count.

Measured with benchmarks/bench_de_overlap.py (Apple M2 Pro, polars 1.x; one
full 10-variant pass over synthetic real/pred DE tables, all outputs verified
bit-identical between old and new):

n_genes=2000, n_sig=100 significant genes/pert:

The old path grows ~quadratically and the new path ~linearly, so the speedup
climbs with n_pert — the axis that matters for the ~18k-perturbation screens
that motivated this. With more significant genes per pert the absolute cost
rises (n_sig=500: 8000 perts = 166.8s old vs 34.0s new), since the irreducible
single-pivot + intersect work grows while the removed redundancy stays fixed.

This is the metric that consumed an entire 8.8h validate-only budget at ~18k
perts (98% of samples in this exact call stack — see issue below); the measured
scaling above is the mechanism behind that, reproduced at tractable sizes.

Reproduce: python benchmarks/bench_de_overlap.py --n-pert 1000 2000 4000 8000

Parity (no behavior change)

Both changes are pure performance; outputs are bit-identical. New
tests/test_de_overlap_equivalence.py:

• compute_overlap matches a from-scratch reference across k ∈ {None,1,2,50,500}
  and metric ∈ {overlap, precision} (exact == on the result dicts).
• the cache collapses the 10 registered variants to one pivot per side
  (len(_top_genes_cache) == 1 after all variants run).
• distinct (sort_by, fdr_threshold) keys stay separate (3 keys → 3 entries).
• the no-significant-genes early-return path still returns all-zeros.

The existing end-to-end tests/test_eval.py (which runs the de/full profile
through the registry, exercising all 10 overlap variants) is unchanged and
still passes.

Fix 3 — DE per-perturbation full-table scans (DENsigCounts, pr/roc AUC)

cell_eval/metrics/_de.py — DENsigCounts.__call__ / compute_generic_auc.

Motivation

With Fix 2 in place, the same ~18k-perturbation OmniPert report run cleared the
overlap pass (the phase that previously stalled 8.6h now completes) and exposed
the next instance of the same per-perturbation-scan pathology downstream.
Faulthandler sampling of the live process (job over a 371k-row DE table) now sat
continuously in:

DENsigCounts.__call__ -> DEResults.get_significant_genes -> polars filter().collect()

Two metrics share the shape — a for pert in iter_perturbations() loop whose
body does a full-table .filter(target == pert):

1. DENsigCounts called get_significant_genes(pert) for the real and
   pred side per perturbation — each a full scan of the whole DE table, so
   ~18k × 2 ≈ 36k scans, where only the significant-gene count is used.
2. compute_generic_auc (backing both pr_auc and roc_auc) built its
   merged frame once (good) but then did merged.filter(target == pert) per
   perturbation — ~18k more full scans, ×2 for the two AUC variants.

Fix

• DENsigCounts: one filter_to_significant().group_by(target).len() per
  side, then reindex over the full iter_perturbations() universe filling 0
  for perts with no significant genes (matching the old empty .size). Only the
  count is consumed downstream, so this is exact.
• compute_generic_auc: one merged.partition_by(target, maintain_order=True, as_dict=True) before the loop; iterate the dict. Perts absent from the map →
  nan (matching the old shape[0] == 0 branch). maintain_order keeps each
  partition in the exact row order the per-pert .filter produced, so the
  labels/scores arrays handed to average_precision_score / roc_curve are
  bit-identical — the per-pert sklearn calls are untouched, only the slicing
  changed from O(perts × rows) to O(rows) total. (partition_by(as_dict=True)
  keys are tuples on newer polars and scalars on older, so they're normalized.)

The audit covered the rest of metrics/_de.py: DESpearmanSignificant,
DESpearmanLFC, DEDirectionMatch, and DESigGenesRecall already use a single
group_by/join (no per-pert loop) and are unchanged.

Parity (no behavior change)

Pure performance; outputs bit-identical. New
tests/test_de_perpert_scan_equivalence.py reproduces the pre-optimization
DENsigCounts and compute_generic_auc verbatim and asserts the new code
matches across a synthetic multi-pert DEComparison — including the
all-significant / all-non-significant perturbations that map to nan, and the
zero-significant reindex path.

Benchmark

benchmarks/bench_de_metrics.py (Apple M2 Pro, polars 1.41; verbatim pre-fix
baselines vs new, all outputs verified bit-identical), n_genes=50 (~the
production DE table's rows-per-pert):

DENsigCounts — the metric the faulthandler actually sat in — was pure
table-scan, so removing it gives 40–172x and the win grows with scale (old is
O(perts × rows), new O(rows)). The pr/roc AUC metrics get a steady ~1.8x: the
per-perturbation sklearn calls and the shared merged build are irreducible,
so only the O(perts × rows) slicing was removed (the absolute time saved still
grows with scale — 0.2s → 1.6s from 1k → 8k perts).

Reproduce: python benchmarks/bench_de_metrics.py --n-pert 1000 2000 4000 8000

Not addressed here (noted for scale)

The anndata-pair metric clustering_agreement remains unprofiled at 18k
perts — a separate potential cliff not yet reached. discrimination_score is
already vectorized (Fix 1). Out of scope for this PR.

Tests / checks

• tests/test_discrimination_score.py: equivalence vs the original loop,
  covering the exclusion, no-exclusion (embedding), exotic-metric fallback,
  duplicate-gene safety net, and target-gene-dominated (degenerate cosine) paths.
• tests/test_de_overlap_equivalence.py: bit-exact overlap/precision
  equivalence + memoization guards (Fix 2).
• tests/test_de_perpert_scan_equivalence.py: bit-exact DENsigCounts /
  pr_auc / roc_auc equivalence vs verbatim pre-optimization references (Fix 3).
• benchmarks/bench_discrimination_score.py, benchmarks/bench_de_overlap.py,
  and benchmarks/bench_de_metrics.py: self-documenting old-vs-new
  microbenchmarks (each keeps a verbatim pre-optimization baseline and asserts
  identical output).

uv run pytest -v          # 89 passed
uv run ruff format --check
uv run ruff check
uv run ty check